import numpy as np
import cv2


def compute_ssd(patch1, patch2):
    """
    计算两个图像块之间的SSD（差值平方和）。

    :param patch1: 图像块1。
    :param patch2: 图像块2。
    :return: SSD值。
    """
    if patch1.shape != patch2.shape:
        return float('inf')  # 返回一个大值表示形状不匹配
    return np.sum((patch1 - patch2) ** 2)


def get_patch(image, center, offset, block_size=3):
    """
    从图像中提取图像块。

    :param image: 输入图像。
    :param center: 图像块的中心位置 (y, x)。
    :param offset: 图像块的偏移量 (dy, dx)。
    :param block_size: 图像块的大小，默认为3。
    :return: 提取的图像块。
    """
    y, x = center
    dy, dx = offset
    half_block = block_size // 2
    y_start = y + dy - half_block
    y_end = y + dy + half_block + 1
    x_start = x + dx - half_block
    x_end = x + dx + half_block + 1

    return image[y_start:y_end, x_start:x_end]


def get_gradient_operator(image, p, direction, block_size=3, step_size=3):
    """
    获取局部区域梯度算子，计算当前像素块与其四个方向邻域像素块的SSD。

    :param image: 输入图像。
    :param p: 当前像素位置 (y, x)。
    :param direction: 方向，四个方向之一。
    :param block_size: 图像块大小，默认为3。
    :param step_size: 偏移的步长，默认为3。
    :return: 局部区域梯度算子。
    """
    offsets = {
        'horizontal': [(0, -step_size), (0, step_size)],                # ——
        'vertical': [(-step_size, 0), (step_size, 0)],                  # |
        'diag1': [(-step_size, -step_size), (step_size, step_size)],    # \
        'diag2': [(-step_size, step_size), (step_size, -step_size)]     # /
    }
    total_ssd = 0
    # 当前像素块
    patch1 = get_patch(image, p, (0, 0), block_size)
    # 该方向偏移的两个像素块
    for offset in offsets[direction]:
        patch2 = get_patch(image, p, offset, block_size)
        total_ssd += compute_ssd(patch1, patch2)

    return total_ssd


def gaussian_kernel(kernel_size=9, sigma=1.5):
    """
    生成高斯核。

    :param kernel_size: 核大小。
    :param sigma: 标准差。
    :return: 高斯核。
    """
    ax = np.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.)
    xx, yy = np.meshgrid(ax, ax)
    kernel = np.exp(-0.5 * (np.square(xx) + np.square(yy)) / np.square(sigma))
    return kernel / np.sum(kernel)


def gaussian_convolution(image, kernel):
    """
    对图像应用高斯卷积。

    :param image: 输入图像。
    :param kernel: 高斯核。
    :return: 卷积后的图像。
    """
    pad_size = kernel.shape[0] // 2
    padded_image = np.pad(image, pad_size, mode='reflect')
    convolved_image = np.zeros_like(image)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            region = padded_image[i:i + kernel.shape[0], j:j + kernel.shape[1]]
            convolved_image[i, j] = np.sum(region * kernel)

    return convolved_image


def adaptive_normalization(rd, beta, chi):
    """
    自适应归一化。

    :param rd: 原始描述符。
    :param beta: 自适应最大阈值。
    :param chi: 自适应最小阈值。
    :return: 归一化后的描述符。
    """
    rd_normalized = np.zeros_like(rd)

    for y in range(rd.shape[0]):
        for x in range(rd.shape[1]):
            min_val = np.min(rd[y, x, :])
            max_val = np.max(rd[y, x, :])
            if min_val > beta:
                rd_normalized[y, x, :] = 1
            elif max_val < chi:
                rd_normalized[y, x, :] = 0
            else:
                if max_val == 0 and min_val == 0:
                    rd_normalized[y, x, :] = 0
                elif max_val == min_val:
                    rd_normalized[y, x, :] = 1
                else:
                    k = -1 / (max_val - min_val)
                    rd_normalized[y, x, :] = k * (rd[y, x, :] - max_val)

    return rd_normalized


def compute_mird(image, k=1.0):
    """
    计算图像的模态独立区域描述符（MIRD）。

    :param image: 输入图像。
    :param k: 调节参数，默认为1.0。
    :return: MIRD 描述符。
    """
    # -----     3.1 多方向区域描述符     ----- #

    #                 ——             |         \        /
    directions = ['horizontal', 'vertical', 'diag1', 'diag2']

    # reference图像的image.shape = (256, 256)
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    h, w = image.shape

    dr = np.zeros((h, w, 4))

    # 计算对应像素块的四个方向上的SSD
    for y in range(4, h - 4):
        for x in range(4, w - 4):
            for i, direction in enumerate(directions):
                dr[y, x, i] = get_gradient_operator(image, (y, x), direction)
    # print('SSD: ', dr)

    # 应用高斯滤波
    gaussian_kernel_ = gaussian_kernel(kernel_size=9, sigma=1.5)
    dr_gaussian = np.zeros_like(dr)
    for i in range(4):
        # 计算距离度量 DR
        dr_gaussian[:, :, i] = gaussian_convolution(dr[:, :, i], gaussian_kernel_)
    # print('DR: ', dr_gaussian)

    # 计算方差度量，v.shape = image.shape
    v = np.mean(dr_gaussian, axis=2)

    # 计算区域描述符 RD
    rd = np.zeros_like(dr)
    for i in range(4):
        # 对于平坦区域，dr_gaussian 和 v 都很小
        # 这种情况下，直接将 rd 设置为 0
        with np.errstate(divide='ignore', invalid='ignore'):
            rd[:, :, i] = np.exp(-dr_gaussian[:, :, i] / v)
        rd[np.isnan(rd)] = 0  # 将可能的 NaN 值设置为 0

    # -----     3.2 归一化自相似区域描述符     ----- #

    # 计算自适应阈值
    sum_rd = np.sum(rd)
    beta = 0.9 * sum_rd / (h * w)
    chi = 0.1 * sum_rd / (h * w)
    # print('rd: ', rd)
    # print('beta = ', beta)
    # print('chi = ', chi)

    # 自适应归一化，计算MIRD
    rd_normalized = adaptive_normalization(rd, beta, chi)

    return rd_normalized
